{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b004705f",
   "metadata": {},
   "source": [
    "# 单动作多智能体-写诗\n",
    "\n",
    "我们可以模拟老师带教学生的模式。来让学生智能体写出一个作品，然后让老师智能体给出改进意见，然后再修改，再给意见，如此反复。现在我们设定，需要多智能体系统为我们根据我们给定的主题提供一篇优美的英文诗，除了完成写作的 agent 外，我们再设定一名精通诗句的老师来查看并修改学生的作品。\n",
    "\n",
    "**流程**  \n",
    "系统首先接收用户的需求（写关于XX主题的诗），在系统中，当学生关注到布置的题目后就会开始创作，当老师发现学生写作完成后就会给学生提出意见，根据老师给出的意见，学生将修改自己的作品，直到设定循环结束。\n",
    "\n",
    "**插入模块**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "af102699",
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "\n",
    "from metagpt.actions import Action, UserRequirement\n",
    "from metagpt.logs import logger\n",
    "from metagpt.roles import Role\n",
    "from metagpt.schema import Message\n",
    "from metagpt.environment import Environment\n",
    "\n",
    "from metagpt.const import MESSAGE_ROUTE_TO_ALL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a9650dcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 声明一个名为 classroom 的 env，我们将所有的 role 都放在其中\n",
    "classroom = Environment()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "268542ae",
   "metadata": {},
   "source": [
    "\n",
    "**定义角色**  \n",
    "定义 Student 角色与 Teacher 角色，与单智能体不同的部分是，我们需要声明每个角色关注的动作（self._watch），只有当关注的动作发生后，角色才会开始行动。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6fcd1ab2",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Student(Role):\n",
    "\n",
    "    name: str = \"xiaoming\"\n",
    "    profile: str = \"Student\"\n",
    "\n",
    "    def __init__(self, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self._init_actions([WritePoem])\n",
    "        self._watch([UserRequirement, ReviewPoem])\n",
    "\n",
    "    async def _act(self) -> Message:\n",
    "        logger.info(f\"{self._setting}: ready to {self.rc.todo}\")\n",
    "        todo = self.rc.todo\n",
    "\n",
    "        msg = self.get_memories()  # 获取所有记忆\n",
    "        # logger.info(msg)\n",
    "        poem_text = await WritePoem().run(msg)\n",
    "        logger.info(f'student : {poem_text}')\n",
    "        msg = Message(content=poem_text, role=self.profile,\n",
    "                      cause_by=type(todo))\n",
    "\n",
    "        return msg\n",
    "\n",
    "class Teacher(Role):\n",
    "\n",
    "    name: str = \"laowang\"\n",
    "    profile: str = \"Teacher\"\n",
    "\n",
    "    def __init__(self, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self._init_actions([ReviewPoem])\n",
    "        self._watch([WritePoem])\n",
    "\n",
    "    async def _act(self) -> Message:\n",
    "        logger.info(f\"{self._setting}: ready to {self.rc.todo}\")\n",
    "        todo = self.rc.todo\n",
    "\n",
    "        msg = self.get_memories()  # 获取所有记忆\n",
    "        poem_text = await ReviewPoem().run(msg)\n",
    "        logger.info(f'teacher : {poem_text}')\n",
    "        msg = Message(content=poem_text, role=self.profile,\n",
    "                      cause_by=type(todo))\n",
    "\n",
    "        return msg"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "443560c7",
   "metadata": {},
   "source": [
    "**定义动作**\n",
    "\n",
    "编写 WritePoem 与 ReviewPoem 方法，在 WritePoem 方法中我们需要实现根据用户提供的主题来编写诗句，并且根据 teacher 的建议修改诗句，在 ReviewPoem 方法中，我们需要读取 student 的诗歌作品，并且给出修改意见。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "728e7fb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class WritePoem(Action):\n",
    "\n",
    "    name: str = \"WritePoem\"\n",
    "\n",
    "    PROMPT_TEMPLATE: str = \"\"\"\n",
    "    Here is the historical conversation record : {msg} .\n",
    "    Write a poem about the subject provided by human, Return only the content of the generated poem with NO other texts.\n",
    "    If the teacher provides suggestions about the poem, revise the student's poem based on the suggestions and return.\n",
    "    your poem:\n",
    "    \"\"\"\n",
    "\n",
    "    async def run(self, msg: str):\n",
    "\n",
    "        prompt = self.PROMPT_TEMPLATE.format(msg = msg)\n",
    "\n",
    "        rsp = await self._aask(prompt)\n",
    "\n",
    "        return rsp\n",
    "\n",
    "class ReviewPoem(Action):\n",
    "\n",
    "    name: str = \"ReviewPoem\"\n",
    "\n",
    "    PROMPT_TEMPLATE: str = \"\"\"\n",
    "\n",
    "    Here is the historical conversation record : {msg} .\n",
    "    Check student-created poems about the subject provided by human and give your suggestions for revisions. You prefer poems with elegant sentences and retro style.\n",
    "    Return only your comments with NO other texts.\n",
    "    your comments:\n",
    "    \"\"\"\n",
    "\n",
    "    async def run(self, msg: str):\n",
    "\n",
    "        prompt = self.PROMPT_TEMPLATE.format(msg = msg)\n",
    "\n",
    "        rsp = await self._aask(prompt)\n",
    "\n",
    "        return rsp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4059ba3",
   "metadata": {},
   "source": [
    "**运行**  \n",
    "提供一个主题，将topic发布在env中运行env，系统就将开始工作了，你可以修改对话轮数（n_round）来达到你希望的效果"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce854f20",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def main(topic: str, n_round=3):\n",
    "\n",
    "    classroom.add_roles([Student(), Teacher()])\n",
    "\n",
    "    classroom.publish_message(\n",
    "        Message(role=\"Human\", content=topic, cause_by=UserRequirement,\n",
    "                send_to='' or MESSAGE_ROUTE_TO_ALL),\n",
    "        peekable=False,\n",
    "    )\n",
    "\n",
    "    while n_round > 0:\n",
    "        # self._save()\n",
    "        n_round -= 1 #如果n_round = 1 ，就只有学生写诗、然后老师没办法进行review\n",
    "        logger.debug(f\"max {n_round=} left.\")\n",
    "\n",
    "        await classroom.run()\n",
    "    return classroom.history\n",
    "\n",
    "await main(topic='wirte a poem about moon')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c69c3ce",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
